// Copyright (c) 2024 Endless Inc.
#include "muxer_mp4_repair.h"
#include "utils_log.h"
#include "utils_file_ex.h"
#include <unistd.h>

namespace El {
namespace Muxer {

bool MP4Repair::Repair(const std::string& inputFile, const std::string& outputFile)
{
    if (!Utils::IsFileExists(inputFile)) {
        EL_ERROR("Input file does not exist: {}", inputFile);
        return false;
    }

    std::string outFile = outputFile;
    if (outFile.empty()) {
        outFile = inputFile + ".tmp";
    }

    // 尝试修复moov atom
    bool success = RepairMoovAtom(inputFile, outFile);
    if (!success) {
        EL_ERROR("Failed to repair MP4 file: {}", inputFile);
        if (Utils::IsFileExists(outFile)) {
            unlink(outFile.c_str());
        }
        return false;
    }

    // 如果输出文件未指定,则覆盖原文件
    if (outputFile.empty()) {
        if (unlink(inputFile.c_str()) != 0 || rename(outFile.c_str(), inputFile.c_str()) != 0) {
            EL_ERROR("Failed to replace original file");
            unlink(outFile.c_str());
            return false;
        }
    }

    return true;
}

bool MP4Repair::RepairMoovAtom(const std::string& inputFile, const std::string& outputFile)
{
    AVFormatContext* inCtx = nullptr;
    AVFormatContext* outCtx = nullptr;

    // 打开输入文件
    if (avformat_open_input(&inCtx, inputFile.c_str(), nullptr, nullptr) < 0) {
        EL_ERROR("Could not open input file: {}", inputFile);
        return false;
    }

    // 读取流信息
    if (avformat_find_stream_info(inCtx, nullptr) < 0) {
        EL_ERROR("Could not find stream info");
        avformat_close_input(&inCtx);
        return false;
    }

    // 创建输出上下文
    if (avformat_alloc_output_context2(&outCtx, nullptr, nullptr, outputFile.c_str()) < 0) {
        EL_ERROR("Could not create output context");
        avformat_close_input(&inCtx);
        return false;
    }

    // 复制流
    for (unsigned int i = 0; i < inCtx->nb_streams; i++) {
        AVStream* inStream = inCtx->streams[i];
        AVStream* outStream = avformat_new_stream(outCtx, nullptr);
        if (!outStream) {
            EL_ERROR("Failed to allocate output stream");
            avformat_free_context(outCtx);
            avformat_close_input(&inCtx);
            return false;
        }

        // 复制编解码器参数
        if (avcodec_parameters_copy(outStream->codecpar, inStream->codecpar) < 0) {
            EL_ERROR("Failed to copy codec params");
            avformat_free_context(outCtx);
            avformat_close_input(&inCtx);
            return false;
        }
        outStream->time_base = inStream->time_base;
    }

    // 打开输出文件
    if (avio_open(&outCtx->pb, outputFile.c_str(), AVIO_FLAG_WRITE) < 0) {
        EL_ERROR("Could not open output file: {}", outputFile);
        avformat_free_context(outCtx);
        avformat_close_input(&inCtx);
        return false;
    }

    // 写入文件头
    if (avformat_write_header(outCtx, nullptr) < 0) {
        EL_ERROR("Error occurred when writing header to output file");
        avio_closep(&outCtx->pb);
        avformat_free_context(outCtx);
        avformat_close_input(&inCtx);
        return false;
    }

    // 复制数据包
    AVPacket packet;
    while (av_read_frame(inCtx, &packet) >= 0) {
        AVStream* inStream = inCtx->streams[packet.stream_index];
        AVStream* outStream = outCtx->streams[packet.stream_index];

        // 转换时间戳
        packet.pts = av_rescale_q_rnd(packet.pts, inStream->time_base, outStream->time_base,
                                    static_cast<AVRounding>(AV_ROUND_NEAR_INF | AV_ROUND_PASS_MINMAX));
        packet.dts = av_rescale_q_rnd(packet.dts, inStream->time_base, outStream->time_base,
                                    static_cast<AVRounding>(AV_ROUND_NEAR_INF | AV_ROUND_PASS_MINMAX));
        packet.duration = av_rescale_q(packet.duration, inStream->time_base, outStream->time_base);
        packet.pos = -1;

        // 写入数据包
        if (av_interleaved_write_frame(outCtx, &packet) < 0) {
            EL_ERROR("Error occurred when writing packet to output file");
            av_packet_unref(&packet);
            avio_closep(&outCtx->pb);
            avformat_free_context(outCtx);
            avformat_close_input(&inCtx);
            return false;
        }
        av_packet_unref(&packet);
    }

    // 写入文件尾
    if (av_write_trailer(outCtx) < 0) {
        EL_ERROR("Error occurred when writing trailer to output file");
        avio_closep(&outCtx->pb);
        avformat_free_context(outCtx);
        avformat_close_input(&inCtx);
        return false;
    }

    // 清理资源
    avio_closep(&outCtx->pb);
    avformat_free_context(outCtx);
    avformat_close_input(&inCtx);

    return true;
}

} // namespace Muxer
} // namespace El 